#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
使用OpenAI API学习谱面创作规律
支持增量式学习方法
支持谱面段落识别和复用
"""

import os
import json
import glob
import re
import numpy as np
from typing import List, Dict, Any, Tuple
from openai import OpenAI
from collections import Counter

# 导入音频分析模块
from beat_analyzer import analyze_audio_features

class ChartSegment:
    """谱面段落类"""
    def __init__(self, notes: List[str], bpm: float = None, metadata: Dict = None):
        self.notes = notes
        self.bpm = bpm
        self.metadata = metadata or {}
        self.frequency = 1
        self.source_charts = []  # 来源谱面
        self.vector = None  # 向量表示
        self.normalized_notes = None  # 标准化后的音符
        
    def __hash__(self):
        # 基于标准化音符内容生成哈希值
        if self.normalized_notes:
            return hash(tuple(self.normalized_notes))
        return hash(tuple(self.notes))
        
    def __eq__(self, other):
        # 比较两个段落是否相同（基于标准化音符）
        if not isinstance(other, ChartSegment):
            return False
        if self.normalized_notes and other.normalized_notes:
            return self.normalized_notes == other.normalized_notes
        return self.notes == other.notes

class ChartLearningEngine:
    def __init__(self, data_dir: str = "data"):
        """
        初始化谱面学习引擎
        
        Args:
            data_dir: 包含谱面JSON文件的目录
        """
        self.data_dir = data_dir
        self.chart_files = []
        self.accumulated_knowledge = ""
        self.start_index = 0
        self.common_segments = {}  # 常见谱面段落 {hash: ChartSegment}
        self.segment_patterns = {}  # 段落模式统计
        
        # 初始化ModelScope客户端
        self.client = OpenAI(
            base_url='https://api-inference.modelscope.cn/v1',
            api_key='ms-4611f4fa-6ec0-4ba7-87da-43032bc42e49'
        )
    
    def normalize_notes(self, notes: List[str]) -> List[str]:
        """
        标准化音符，消除位置平移影响
        
        Args:
            notes: 音符列表
            
        Returns:
            标准化后的音符列表
        """
        if not notes:
            return notes
            
        normalized = []
        # 解析所有音符，找出最小按键位置
        all_positions = []
        parsed_lines = []
        
        for line in notes:
            if line.startswith('{') and '},' in line:
                try:
                    # 提取note值
                    note_value_str = line[1:line.index('}')]
                    note_value = int(note_value_str)
                    # 提取音符部分
                    notes_part = line[line.index('},') + 2:]
                    note_list = notes_part.split(',') if notes_part else []
                    parsed_lines.append((note_value, note_list))
                    
                    # 收集所有按键位置
                    for note in note_list:
                        if note and note[0].isdigit():
                            pos = int(note[0])
                            all_positions.append(pos)
                except:
                    parsed_lines.append(line)
            else:
                parsed_lines.append(line)
        
        if not all_positions:
            return notes
            
        # 计算位置偏移（将最小位置移到1）
        min_pos = min(all_positions)
        offset = 1 - min_pos
        
        # 应用偏移
        for item in parsed_lines:
            if isinstance(item, tuple):
                note_value, note_list = item
                new_notes = []
                for note in note_list:
                    # 处理各种类型的音符
                    if note and note[0].isdigit():
                        pos = int(note[0])
                        new_pos = pos + offset
                        # 确保位置在1-8范围内
                        if 1 <= new_pos <= 8:
                            new_note = str(new_pos) + note[1:]
                        else:
                            # 如果超出范围，保持原样或做适当处理
                            new_note = note
                        new_notes.append(new_note)
                    else:
                        new_notes.append(note)
                
                # 重新构建行
                notes_str = ','.join(new_notes)
                normalized.append(f"{{{note_value}}},{notes_str},")
            else:
                normalized.append(item)
                
        return normalized
    
    def vectorize_segment(self, notes: List[str]) -> np.ndarray:
        """
        将谱面段落转换为向量表示
        
        Args:
            notes: 音符列表
            
        Returns:
            向量表示
        """
        # 简化的向量化方法：统计各种音符类型的数量
        note_stats = {
            'tap': 0,
            'hold': 0,
            'slide': 0,
            'break': 0,
            'multi': 0,
            'total': len(notes)
        }
        
        for line in notes:
            if line.startswith('{') and '},' in line:
                try:
                    notes_part = line[line.index('},') + 2:]
                    note_list = notes_part.split(',') if notes_part else []
                    for note in note_list:
                        if not note:
                            continue
                        if 'h[' in note:
                            note_stats['hold'] += 1
                        elif 's' in note or '>' in note or '<' in note or '-' in note:
                            note_stats['slide'] += 1
                        elif 'b' in note:
                            note_stats['break'] += 1
                        elif '/' in note:
                            note_stats['multi'] += 1
                        elif note[0].isdigit():
                            note_stats['tap'] += 1
                except:
                    pass
        
        # 转换为向量
        vector = np.array([
            note_stats['tap'],
            note_stats['hold'],
            note_stats['slide'],
            note_stats['break'],
            note_stats['multi'],
            note_stats['total']
        ], dtype=float)
        
        # 归一化
        if note_stats['total'] > 0:
            vector[:-1] /= note_stats['total']
            
        return vector
    
    def segments_similar(self, seg1: ChartSegment, seg2: ChartSegment, threshold: float = 0.9) -> bool:
        """
        判断两个谱面段落是否相似
        
        Args:
            seg1: 第一个段落
            seg2: 第二个段落
            threshold: 相似度阈值
            
        Returns:
            是否相似
        """
        # 首先检查标准化音符
        if seg1.normalized_notes and seg2.normalized_notes:
            if seg1.normalized_notes == seg2.normalized_notes:
                return True
        
        # 检查向量相似度
        if seg1.vector is not None and seg2.vector is not None:
            # 计算余弦相似度
            dot_product = np.dot(seg1.vector, seg2.vector)
            norm1 = np.linalg.norm(seg1.vector)
            norm2 = np.linalg.norm(seg2.vector)
            
            if norm1 > 0 and norm2 > 0:
                cosine_sim = dot_product / (norm1 * norm2)
                return cosine_sim >= threshold
                
        return False
    
    def load_progress(self) -> bool:
        """
        从进度文件加载学习进度
        
        Returns:
            是否成功加载进度
        """
        progress_file = "learning_progress.json"
        if not os.path.exists(progress_file):
            print("未找到进度文件，将从头开始学习")
            return False
        
        try:
            with open(progress_file, 'r', encoding='utf-8') as f:
                progress_data = json.load(f)
            
            self.accumulated_knowledge = progress_data.get("accumulated_knowledge", "")
            self.start_index = progress_data.get("current_index", 0)
            # 加载常见段落
            if "common_segments" in progress_data:
                self.common_segments = {}
                for seg_hash, seg_data in progress_data["common_segments"].items():
                    segment = ChartSegment(seg_data["notes"], seg_data["bpm"], seg_data["metadata"])
                    segment.frequency = seg_data["frequency"]
                    segment.source_charts = seg_data["source_charts"]
                    segment.normalized_notes = seg_data.get("normalized_notes")
                    vector_data = seg_data.get("vector")
                    if vector_data:
                        segment.vector = np.array(vector_data)
                    self.common_segments[seg_hash] = segment
            
            print(f"已从进度文件加载学习进度: {self.start_index}")
            if self.accumulated_knowledge:
                print(f"已加载的知识摘要: {self.accumulated_knowledge[:200]}...")
            return True
        except Exception as e:
            print(f"加载进度文件时出错: {e}")
            return False
    
    def scan_chart_files(self):
        """扫描data目录下的所有谱面JSON文件"""
        pattern = os.path.join(self.data_dir, "*.json")
        self.chart_files = glob.glob(pattern)
        self.chart_files.sort()  # 按文件名排序
        
        print(f"总共找到 {len(self.chart_files)} 个谱面文件")
        return self.chart_files
    
    def read_chart_file(self, chart_file: str) -> Dict[str, Any]:
        """
        读取谱面JSON文件内容
        
        Args:
            chart_file: 谱面JSON文件路径
            
        Returns:
            谱面文件内容（JSON对象）
        """
        try:
            with open(chart_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"读取谱面文件 {chart_file} 时出错: {e}")
            return {}
    
    def extract_chart_content(self, chart_data: Dict[str, Any], difficulty: int = 2) -> str:
        """
        从JSON数据中提取指定难度的谱面内容
        
        Args:
            chart_data: 谱面JSON数据
            difficulty: 要提取的谱面难度等级
            
        Returns:
            谱面内容字符串
        """
        # 在charts数组中查找指定难度的谱面
        for chart in chart_data.get("chart_info", {}).get("charts", []):
            if chart.get("difficulty") == difficulty:
                return chart.get("chart", "")
        
        # 如果没找到指定难度，返回空字符串
        return ""
    
    def extract_chart_segments(self, chart_content: str, min_segment_length: int = 8, max_segment_length: int = 32) -> List[ChartSegment]:
        """
        从谱面内容中提取段落
        
        Args:
            chart_content: 谱面内容
            min_segment_length: 最小段落长度
            max_segment_length: 最大段落长度
            
        Returns:
            谱面段落列表
        """
        segments = []
        lines = chart_content.split('\n')
        
        # 提取谱面行（去除元数据和结束标记）
        chart_lines = []
        bpm = None
        
        for line in lines:
            line = line.strip()
            if line.startswith('&wholebpm='):
                try:
                    bpm = float(line.split('=')[1])
                except:
                    bpm = 120.0
            elif line.startswith('{') and '},' in line:
                chart_lines.append(line)
        
        # 生成不同长度的段落
        for length in range(min_segment_length, min(max_segment_length + 1, len(chart_lines) + 1)):
            for i in range(len(chart_lines) - length + 1):
                segment_lines = chart_lines[i:i + length]
                segment = ChartSegment(segment_lines, bpm)
                # 添加标准化音符和向量表示
                segment.normalized_notes = self.normalize_notes(segment_lines)
                segment.vector = self.vectorize_segment(segment_lines)
                segments.append(segment)
                
        return segments
    
    def identify_common_segments(self, new_segments: List[ChartSegment], chart_name: str):
        """
        识别常见谱面段落
        
        Args:
            new_segments: 新的谱面段落
            chart_name: 谱面名称
        """
        for segment in new_segments:
            seg_hash = hash(segment)
            found_similar = False
            
            # 首先检查完全相同的段落
            if seg_hash in self.common_segments:
                # 段落已存在，增加频率
                self.common_segments[seg_hash].frequency += 1
                if chart_name not in self.common_segments[seg_hash].source_charts:
                    self.common_segments[seg_hash].source_charts.append(chart_name)
                found_similar = True
            else:
                # 检查相似的段落
                for existing_hash, existing_segment in self.common_segments.items():
                    if self.segments_similar(segment, existing_segment):
                        # 找到相似段落，增加频率
                        existing_segment.frequency += 1
                        if chart_name not in existing_segment.source_charts:
                            existing_segment.source_charts.append(chart_name)
                        found_similar = True
                        break
            
            # 如果没有找到相似段落，则添加新段落
            if not found_similar:
                segment.source_charts.append(chart_name)
                self.common_segments[seg_hash] = segment
    
    def analyze_audio_for_chart(self, chart_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        从谱面数据中获取音频特征
        
        Args:
            chart_data: 谱面JSON数据
            
        Returns:
            音频特征数据
        """
        # 直接从JSON数据中获取音频特征
        audio_features = chart_data.get("audio_features", {})
        if audio_features:
            print(f"已从JSON数据中获取音频特征")
            return audio_features
        else:
            print("JSON数据中未包含音频特征")
            return None
    
    def learn_single_chart(self, chart_data: Dict[str, Any]) -> str:
        """
        学习单个谱面并与已有知识整合
        
        Args:
            chart_data: 单个谱面数据（JSON对象）
            
        Returns:
            更新后的知识
        """
        chart_file = chart_data['file']
        chart_json = chart_data['content']
        print(f"正在学习谱面: {chart_file}")
        
        # 提取谱面内容
        chart_content = self.extract_chart_content(chart_json, difficulty=2)
        # 获取音频特征
        audio_features = self.analyze_audio_for_chart(chart_json)
        
        # 提取谱面段落
        if chart_content:
            segments = self.extract_chart_segments(chart_content)
            chart_name = os.path.basename(chart_file)
            self.identify_common_segments(segments, chart_name)
        
        # 读取谱面生成代码
        generator_code = ""
        generator_file_path = os.path.join(os.path.dirname(__file__), "generate.py")
        if os.path.exists(generator_file_path):
            try:
                with open(generator_file_path, 'r', encoding='utf-8') as f:
                    generator_code = f.read()
            except Exception as e:
                print(f"读取谱面生成代码时出错: {e}")
        
        # 构造提示词
        prompt = f"""
你是一位音乐游戏谱面分析专家。{'这是我们已经积累的谱面知识：' + self.accumulated_knowledge if self.accumulated_knowledge else '请分析以下谱面并总结其规律'}

请分析以下谱面文件，并{'将其规律与已有知识整合' if self.accumulated_knowledge else '总结其规律'}：

谱面信息：
文件: {chart_file}
标题: {chart_json.get("chart_info", {}).get("title", "未知")}
艺术家: {chart_json.get("chart_info", {}).get("artist", "未知")}
"""
        
        if audio_features:
            prompt += f"音频特征: {json.dumps(audio_features['metadata'], ensure_ascii=False, indent=2)}\n"
        
        if chart_content:
            prompt += f"等级2谱面数据:\n{chart_content}\n"
        
        if generator_code:
            prompt += f"\n当前的谱面生成代码:\n{generator_code}\n"
        
        prompt += f"""
请根据以上信息，分析该谱面的设计规律，并针对谱面生成代码提出具体的修改建议：

1. 分析谱面结构组织方式和音符类型使用规律
2. 分析谱面难度与参数设置的关系
3. 分析音符分布与音乐节奏的对应关系
4. 识别可复用的谱面段落模式
5. 针对谱面生成代码提出具体的修改建议，包括：
   - 需要增加或修改的函数/类
   - 需要调整的参数或算法
   - 可以优化的代码结构
   - 具体的代码实现建议

请以清晰、结构化的方式输出分析结果和修改建议。
"""

        try:
            response = self.client.chat.completions.create(
                model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                messages=[
                    {
                        'role': 'system',
                        'content': '你是一位专业的音乐游戏谱面分析专家和Python程序员，能够从谱面数据和音频特征中发现深层规律，并能针对代码提出具体的修改建议。'
                    },
                    {
                        'role': 'user',
                        'content': prompt
                    }
                ],
                temperature=0.7,
                max_tokens=4000
            )
            
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用OpenAI API时出错: {e}")
            return ""
    
    def incremental_learning(self, sample_count: int = None, resume: bool = True):
        """
        增量式学习谱面规律
        
        Args:
            sample_count: 采样谱面数量，如果为None则学习所有谱面
            resume: 是否从保存的进度继续学习
        """
        # 扫描谱面文件
        self.scan_chart_files()
        
        if not self.chart_files:
            print("未找到任何谱面文件")
            return
        
        # 尝试从进度文件恢复
        if resume and self.load_progress():
            print(f"从第 {self.start_index + 1} 个谱面继续学习")
            start_from = self.start_index
        else:
            start_from = 0
        
        # 确定要处理的谱面
        process_files = self.chart_files[start_from:]
        if sample_count is not None:
            process_files = process_files[:sample_count] if len(process_files) > sample_count else process_files
            
        total_to_process = len(process_files)
        print(f"将处理 {total_to_process} 个谱面")
        
        if total_to_process == 0:
            print("没有需要处理的谱面")
            return
        
        # 逐个学习谱面
        for i, chart_file in enumerate(process_files):
            actual_index = start_from + i
            print(f"\n处理进度: {actual_index+1}/{len(self.chart_files)} (当前批次: {i+1}/{total_to_process})")
            
            # 读取谱面数据
            content = self.read_chart_file(chart_file)
            if not content:
                continue
                
            chart_data = {
                'file': chart_file,
                'content': content
            }
            
            # 学习单个谱面
            new_knowledge = self.learn_single_chart(chart_data)
            
            if new_knowledge:
                # 更新累积知识
                if self.accumulated_knowledge:
                    # 整合新旧知识
                    integration_prompt = f"""
你是一位音乐游戏谱面分析专家。请整合以下两部分知识：

1. 已有知识：
{self.accumulated_knowledge}

2. 新增知识（来自谱面 {chart_file}）：
{new_knowledge}

此外，这是当前的谱面生成代码：
"""

                    # 读取谱面生成代码
                    generator_code = ""
                    generator_file_path = os.path.join(os.path.dirname(__file__), "generate.py")
                    if os.path.exists(generator_file_path):
                        try:
                            with open(generator_file_path, 'r', encoding='utf-8') as f:
                                generator_code = f.read()
                        except Exception as e:
                            print(f"读取谱面生成代码时出错: {e}")
                    
                    integration_prompt += f"{generator_code}\n"
                    
                    integration_prompt += """
请将这两部分知识整合，形成更全面、更深入的谱面设计规律总结。
注意去重和补充，使知识体系更加完善。
同时，请针对谱面生成代码提出具体的修改建议，包括：
1. 需要增加或修改的函数/类
2. 需要调整的参数或算法
3. 可以优化的代码结构
4. 具体的代码实现建议

请以清晰、结构化的方式输出整合后的知识和修改建议。
"""

                    try:
                        response = self.client.chat.completions.create(
                            model='Qwen/Qwen3-Next-80B-A3B-Instruct',
                            messages=[
                                {
                                    'role': 'system',
                                    'content': '你是一位专业的音乐游戏谱面分析专家和Python程序员，擅长知识整合与总结，能够针对代码提出具体的修改建议。'
                                },
                                {
                                    'role': 'user',
                                    'content': integration_prompt
                                }
                            ],
                            temperature=0.7,
                            max_tokens=4000
                        )
                        
                        self.accumulated_knowledge = response.choices[0].message.content
                    except Exception as e:
                        print(f"整合知识时出错: {e}")
                        # 如果整合失败，简单拼接新知识
                        self.accumulated_knowledge += f"\n\n从谱面 {chart_file} 中学到的新规律:\n{new_knowledge}"
                else:
                    # 第一个谱面的知识
                    self.accumulated_knowledge = new_knowledge
            
            print(f"已完成谱面 {chart_file} 的学习和知识整合")
            
            # 每处理5个谱面保存一次进度
            if (i + 1) % 5 == 0 or (i + 1) == total_to_process:
                self._save_progress(actual_index + 1, len(self.chart_files))
        
        # 保存最终结果
        if self.accumulated_knowledge:
            output_file = "incremental_learning_result.txt"
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(self.accumulated_knowledge)
            
            # 保存常见段落
            self._save_common_segments()
            
            print(f"\n增量学习完成，最终结果已保存到 {output_file}")
            print("\n最终知识摘要:")
            print(self.accumulated_knowledge[:500] + "..." if len(self.accumulated_knowledge) > 500 else self.accumulated_knowledge)
        else:
            print("学习失败，未生成任何知识")
    
    def _save_common_segments(self):
        """保存常见谱面段落"""
        if not self.common_segments:
            return
            
        # 按频率排序，只保存高频段落
        sorted_segments = sorted(self.common_segments.items(), 
                               key=lambda x: x[1].frequency, 
                               reverse=True)
        
        # 只保存频率大于1的段落
        frequent_segments = {k: v for k, v in sorted_segments if v.frequency > 1}
        
        segments_data = {}
        for seg_hash, segment in frequent_segments.items():
            segments_data[seg_hash] = {
                "notes": segment.notes,
                "bpm": segment.bpm,
                "metadata": segment.metadata,
                "frequency": segment.frequency,
                "source_charts": segment.source_charts,
                "normalized_notes": segment.normalized_notes,
                "vector": segment.vector.tolist() if segment.vector is not None else None
            }
        
        segments_file = "common_segments.json"
        try:
            with open(segments_file, 'w', encoding='utf-8') as f:
                json.dump(segments_data, f, ensure_ascii=False, indent=2)
            print(f"常见谱面段落已保存到 {segments_file}")
            print(f"共保存 {len(frequent_segments)} 个高频段落")
        except Exception as e:
            print(f"保存常见段落时出错: {e}")
    
    def _save_progress(self, current_index: int, total_count: int):
        """
        保存学习进度
        
        Args:
            current_index: 当前处理的谱面索引（从1开始）
            total_count: 总谱面数量
        """
        progress_data = {
            "current_index": current_index,
            "total_count": total_count,
            "accumulated_knowledge": self.accumulated_knowledge,
            "processed_files": self.chart_files[:current_index],
            "common_segments": {}
        }
        
        # 保存常见段落
        for seg_hash, segment in self.common_segments.items():
            seg_data = {
                "notes": segment.notes,
                "bpm": segment.bpm,
                "metadata": segment.metadata,
                "frequency": segment.frequency,
                "source_charts": segment.source_charts,
                "normalized_notes": segment.normalized_notes
            }
            if segment.vector is not None:
                seg_data["vector"] = segment.vector.tolist()
            progress_data["common_segments"][seg_hash] = seg_data
        
        progress_file = "learning_progress.json"
        try:
            with open(progress_file, 'w', encoding='utf-8') as f:
                json.dump(progress_data, f, ensure_ascii=False, indent=2)
            print(f"进度已保存: {current_index}/{total_count}")
        except Exception as e:
            print(f"保存进度时出错: {e}")

def main():
    """主函数"""
    engine = ChartLearningEngine()
    
    print("开始使用AI进行增量式谱面学习")
    print("="*50)
    
    # 进行增量学习，支持从进度恢复
    engine.incremental_learning(sample_count=3, resume=True)  # 只处理3个文件用于测试

if __name__ == "__main__":
    print("启动脚本...")
    main()
    print("脚本执行完成")